#include "hybrid.h"

int check_pubkey_file()
{
  // 检查公钥文件是否存在
  #ifdef _WIN32
    if((0 == _access("./KEY/kyberPubKey.bin",F_OK)) && (0 == access("./KEY/SMPubKey.bin",F_OK)))
      return 1;
    else
      return 0;
  #else
    if((0 == access("./KEY/kyberPubKey.bin",F_OK)) && (0 == access("./KEY/SMPubKey.bin",F_OK)))
      return 1;
    else
      return 0;
  #endif
}

int check_prikey_file()
{
  // 检查私钥文件是否存在
  #ifdef _WIN32
    if((0 == _access("./KEY/kyberPriKey.bin",0)) && (0 == access("./KEY/SMPriKey.bin",F_OK)))
      return 1;
    else
      return 0;
  #else
    if((0 == access("./KEY/kyberPriKey.bin",F_OK)) && (0 == access("./KEY/SMPriKey.bin",F_OK)))
      return 1;
    else
      return 0;
  #endif
}

int gene_key()
{
    int tmp;

    // Kyber
    uint8_t pk_gen[KYBER_INDCPA_PUBLICKEYBYTES];  //生成的公钥
    uint8_t sk_gen[KYBER_INDCPA_PUBLICKEYBYTES];  //生成的私钥

    FILE* outKyberPriFile;  //保存私钥的文件指针
    FILE* outKyberPubFile;  //保存公钥的文件指针

    //Generete Kyber Key
    indcpa_keypair(pk_gen, sk_gen);

    // save secret key
    if (!(outKyberPriFile = fopen("./KEY/kyberPriKey.bin", "wb")))
    {
      return -1;
    }
    fwrite(sk_gen, sizeof(char), KYBER_INDCPA_SECRETKEYBYTES, outKyberPriFile);
    fclose(outKyberPriFile);
    printf("Save Kyber SK Success\n");

    // save public key
    if (!(outKyberPubFile = fopen("./KEY/kyberPubKey.bin", "wb")))
    {
      return -1;
    }
    fwrite(pk_gen, sizeof(char), KYBER_INDCPA_PUBLICKEYBYTES, outKyberPubFile);
    fclose(outKyberPubFile);
    printf("Save Kyber PK Success\n");
    // End Kyber

    // SM2
    SM2_KEY sm2_key;  //SM2公私钥结构
    FILE* outSMPriFile; //保存私钥的文件指针
    FILE* outSMPubFile; //保存公钥的文件指针

    sm2_key_generate(&sm2_key);

    // save secret key
    if (!(outSMPriFile = fopen("./KEY/SMPriKey.bin", "wb")))
    {
      return -1;
    }
    fwrite(sm2_key.private_key, sizeof(char), 32, outSMPriFile);
    fclose(outSMPriFile);
    printf("Save SM2 SK Success\n");

    // save public key
    if (!(outSMPubFile = fopen("./KEY/SMPubKey.bin", "wb")))
    {
      return -1;
    }
    fwrite(&sm2_key.public_key, sizeof(SM2_POINT), 1, outSMPubFile);
    fclose(outSMPubFile);
    printf("Save SM2 PK Success\n");
    // End SM2
  
  return 1;
}

int hybrid_encrypt(unsigned char* message, unsigned char* cipherText, int* cipherLen)
{   
  if(0==check_pubkey_file())
    return -1;
    
  unsigned char msgKyber[ENCRYPT_KEY_LEN];  //Kyber要处理的原始内容
  unsigned char msgSM[ENCRYPT_KEY_LEN]; //SM2要处理的原始内容
  FILE* inKyberPubFile; //Kyber读取公钥的文件指针
  FILE* inSMPubFile;  //SM2读取公钥的文件指针

  // Kyber
  unsigned char kyberPubKey[KYBER_INDCPA_PUBLICKEYBYTES]; //Kyber的公钥
  unsigned char coins[KYBER_SYMBYTES];  //Kyber的随机数

  memcpy(msgKyber, message, ENCRYPT_KEY_LEN);
  // printf("原始密钥Kyber:");
  // for (int i = 0; i < ENCRYPT_KEY_LEN; i++)
  // {
  //   printf("%02X", msgKyber[i]);
  // }

  // read public key
  if (!(inKyberPubFile = fopen("./KEY/kyberPubKey.bin", "rb")))
    return -1;
  fread(kyberPubKey, sizeof(char), KYBER_INDCPA_PUBLICKEYBYTES, inKyberPubFile);
  fclose(inKyberPubFile);

  kyber_randombytes(coins, KYBER_SYMBYTES);

  indcpa_enc(cipherText, msgKyber, kyberPubKey, coins);
  // End Kyber

  // SM2
  SM2_KEY SMPubKey; //SM2的公钥
  size_t SMLen; //SM2加密后密文长度

  if (!(inSMPubFile = fopen("./KEY/SMPubKey.bin", "rb")))
    return -1;
  fread(&SMPubKey.public_key, sizeof(SM2_POINT), 1, inSMPubFile);
  fclose(inSMPubFile);

  memcpy(msgSM, message+ENCRYPT_KEY_LEN, ENCRYPT_KEY_LEN);

  // printf("\n原始密钥SM2:");
  // for (int i = 0; i < ENCRYPT_KEY_LEN; i++)
  // {
  //   printf("%02X", msgSM[i]);
  // }
  // printf("\n");

  sm2_encrypt(&SMPubKey, msgSM, ENCRYPT_KEY_LEN, cipherText+KYBER_INDCPA_BYTES, &SMLen);
  // End SM2

  *cipherLen = KYBER_INDCPA_BYTES + SMLen;  //传回密文总长度

  return 1;
}

int hybrid_decrypt(unsigned char* cipherText, int cipherLen, unsigned char* message)
{
  if(0==check_prikey_file())
    return -1;
      
  unsigned char ctKyber[KYBER_INDCPA_BYTES]={0};  //Kyber要处理的密文
  unsigned char ctSM[SM2_CIPHER_LEN]={0}; //SM2要处理的密文
  FILE* inKyberPriFile; //Kyber读取私钥的文件指针
  FILE* inSMPriFile;  //SM2读取私钥的文件指针

  // Kyber
  unsigned char kyberPriKey[KYBER_INDCPA_PUBLICKEYBYTES]; //Kyber的私钥

  memcpy(ctKyber, cipherText, KYBER_INDCPA_BYTES);

  // read public key
  if (!(inKyberPriFile = fopen("./KEY/kyberPriKey.bin", "rb")))
    return -1;
  fread(kyberPriKey, sizeof(char), KYBER_INDCPA_SECRETKEYBYTES, inKyberPriFile);
  fclose(inKyberPriFile);

  indcpa_dec(message, ctKyber, kyberPriKey);
  // End Kyber

  // SM2
  SM2_KEY SMPriKey; //SM2的私钥
  size_t SMInLen; //SM2密文长度
  unsigned char ptSM[32]; //SM解密后的明文信息
  size_t SMOutLen;  //SM2解密后明文长度

  if (!(inSMPriFile = fopen("./KEY/SMPriKey.bin", "rb")))
    return -1;
  fread(SMPriKey.private_key, sizeof(char), 32, inSMPriFile);
  fclose(inSMPriFile);

  memcpy(ctSM, cipherText+KYBER_INDCPA_BYTES, SM2_CIPHER_LEN);
  SMInLen = cipherLen - KYBER_INDCPA_BYTES;
  sm2_decrypt(&SMPriKey, ctSM, SMInLen, ptSM, &SMOutLen);

  memcpy(message+ENCRYPT_KEY_LEN, ptSM, ENCRYPT_KEY_LEN);
  // End SM2

  return 1;
}

int cal_sm3(unsigned char* key, unsigned char* sm3hash)
{
  SM3_CTX sm3_ctx;

	sm3_init(&sm3_ctx);
	sm3_update(&sm3_ctx, key, ORIGIN_KEY_LEN);
	sm3_finish(&sm3_ctx, sm3hash);
}

int main()
{
  unsigned char originKey[ORIGIN_KEY_LEN]={0};  //原始密钥，长度64
  unsigned char originHash[HASH_LEN];
  unsigned char cipherKey[KYBER_INDCPA_BYTES + SM2_CIPHER_LEN]={0}; //密文
  int cipherLen=0;  //密文长度
  unsigned char unpackKey[ORIGIN_KEY_LEN]={0};  //解密后密钥
  unsigned char unpackHash[HASH_LEN];

  // gene_key();

  kyber_randombytes(originKey, ORIGIN_KEY_LEN);

  printf("原始密钥(Kyber||SM2):\n");
  for (int i = 0; i < ORIGIN_KEY_LEN; i++)
  {
    printf("%02X", originKey[i]);
    if(i==32)
      printf(" ");
  }
  printf("\n");

  cal_sm3(originKey, originHash);
  printf("原始密钥Hash:\n");
  for (int i = 0; i < HASH_LEN; i++)
  {
    printf("%02X", originHash[i]);
  }
  printf("\n");

  hybrid_encrypt(originKey, cipherKey, &cipherLen);

  printf("加密后密钥(部分):");
  for (int i = 0; i < ORIGIN_KEY_LEN; i++)
  {
    printf("%02X", cipherKey[i]);
  }
  printf("\n");

  printf("加密密钥长度:%d\n", cipherLen);

  hybrid_decrypt(cipherKey, cipherLen, unpackKey);

  printf("解密密钥(Kyber||SM2):\n");
  for (int i = 0; i < ORIGIN_KEY_LEN; i++)
  {
    printf("%02X", unpackKey[i]);
    if(i==32)
      printf(" ");
  }
  printf("\n");

  cal_sm3(unpackKey, unpackHash);
  printf("解密密钥Hash:\n");
  for (int i = 0; i < HASH_LEN; i++)
  {
    printf("%02X", unpackHash[i]);
  }
  printf("\n");

  if (memcmp(originHash, unpackHash, HASH_LEN) != 0)
    printf("WRONG!\n");
  else
    printf("Right!!!\n");

  return 1;
}